/*
 * Copyright 2025 The Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 *@file    : reprojection_factor.h
 *@brief   : reprojection factor
 *@authors : zhanglei
 *@version : v1.0
 *@data    : 2025/9/21
 *
 */
#ifndef REPROJECTION_FACTOR_H_
#define REPROJECTION_FACTOR_H_

#include <ceres/ceres.h>
#include "common/type.h"
#include "feature/feature.h"

namespace optimizer {

using namespace feature;

struct ReprojectionFactor : public ceres::SizedCostFunction<2, 1, 2, 3> {
    ReprojectionFactor(const Vector2d& point, const Eigen::Matrix3d& intrinsic_matrix, const Eigen::Matrix4d& extrinsic_matrix) : undistorted_point_(point), intrinsic_matrix_(intrinsic_matrix), extrinsic_matrix_(extrinsic_matrix) {}
    
    inline Eigen::Vector2d CamProject(const Eigen::Vector3d& pos_c) const {
        const float& fx = intrinsic_matrix_(0, 0);
        const float& fy = intrinsic_matrix_(1, 1);
        const float& cx = intrinsic_matrix_(0, 2);
        const float& cy = intrinsic_matrix_(1, 2);
        return {fx * pos_c(0) / pos_c(2) + cx, fy * pos_c(1) / pos_c(2) + cy};
    }

    inline Eigen::Vector2d NormalizedToImagePoint(const Eigen::Vector2d& normalized_pt) const {
        const float& fx = intrinsic_matrix_(0, 0);
        const float& fy = intrinsic_matrix_(1, 1);
        const float& cx = intrinsic_matrix_(0, 2);
        const float& cy = intrinsic_matrix_(1, 2);
        return {fx * normalized_pt(0) + cx, fy * normalized_pt(1) + cy};
    }

    Eigen::Vector2d EvaluateNormalizedResidual(const Vector2d& normalized_pt, const Eigen::Vector3d& world_point, const double& theta, const Vector2d& translation) {
        Eigen::Matrix4d Twb = SE2::matrix4d(SO2(theta), translation);
        Eigen::Matrix4d Twc = Twb * extrinsic_matrix_;
        Eigen::Matrix4d Tcw = Twc.inverse();
        Eigen::Vector3d lw = world_point;
        Eigen::Matrix3d Rcw = Tcw.block<3,3>(0,0);
        Eigen::Vector3d tcw = Tcw.block<3,1>(0,3);
        Eigen::Vector3d lc = Rcw * lw + tcw;
        Eigen::Vector2d cam_project = lc.head<2>()/lc[2];
        Eigen::Vector2d residual = cam_project - normalized_pt; 

        // LOG(INFO) << " normalized residual ===== : " << residual.transpose() << std::endl;
        // // test
        // Eigen::Vector2d obs = NormalizedToImagePoint(normalized_pt);
        // Eigen::Vector2d predict = NormalizedToImagePoint(cam_project);

        // LOG(INFO) << " obs ===== : " << obs.transpose() << std::endl;
        // LOG(INFO) << " predict ===== : " << predict.transpose() << std::endl;

        return residual;
    }

    Eigen::Vector2d EvaluateResidual(const Eigen::Vector3d& world_point, const double& theta, const Vector2d& translation) {
        Eigen::Matrix4d Twb = SE2::matrix4d(SO2(theta), translation);
        Eigen::Matrix4d Twc = Twb * extrinsic_matrix_;
        Eigen::Matrix4d Tcw = Twc.inverse();
        Eigen::Vector3d lw = world_point;
        Eigen::Matrix3d Rcw = Tcw.block<3,3>(0,0);
        Eigen::Vector3d tcw = Tcw.block<3,1>(0,3);
        Eigen::Vector3d lc = Rcw * lw + tcw;

        Eigen::Vector2d residual = CamProject(lc) - Eigen::Vector2d(undistorted_point_); 

        // LOG(INFO) << " intrinsic_matrix_ ===== : " << intrinsic_matrix_ << std::endl;
        
        // LOG(INFO) << " \nundistorted_point_ ===== : " << undistorted_point_.transpose() << std::endl;
        // LOG(INFO) << " CamProject(lc) ===== : " << CamProject(lc).transpose() << std::endl;

        // LOG(INFO) << " residual ===== : " << residual.transpose() << std::endl;

        return residual;
    }

    virtual bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const {
        double theta_i = parameters[0][0];
        Eigen::Vector2d pose_i(parameters[1][0], parameters[1][1]);
        Eigen::Vector3d world_point(parameters[2][0], parameters[2][1], parameters[2][2]);

        Eigen::Matrix4d Twb = SE2::matrix4d(SO2(theta_i), pose_i);

        // LOG(INFO) << " Twb ===== : " << Twb << std::endl;

        Eigen::Matrix4d Twc = Twb * extrinsic_matrix_;
        Eigen::Matrix4d Tcw = Twc.inverse();

        const Vector2d& undist_keypt = undistorted_point_; 
        
        // LOG(INFO) << " world_point ===== : " << world_point.transpose() << std::endl;

        //best odometry 1.0e-3, 1.0e3
        float Sigma_rotxy = 1.0e-10; // 1.0e-3
        // best 5.0
        float Sigma_z = 5.0;
        const float level_sigma_sq = 1.0;
          
        // compute covariance/information
        Eigen::Matrix2d Sigma_u = Eigen::Matrix2d::Identity() * level_sigma_sq;
        Eigen::Vector3d lw = world_point;
        Eigen::Matrix3d Rcw = Tcw.block<3,3>(0,0);
        Eigen::Vector3d tcw = Tcw.block<3,1>(0,3);

        // LOG(INFO) << " world_point ===== : " << world_point.transpose() << std::endl;

        // //std::cout << " lw ===== : " << lw << std::endl;
            
        Eigen::Vector3d lc = Rcw * lw + tcw;
        double zc = lc(2);
        double zc_inv = 1.0 / zc;
        double zc_inv2 = zc_inv * zc_inv;

        const float& fx = intrinsic_matrix_(0, 0);
        const float& fy = intrinsic_matrix_(1, 1);

        Eigen::Matrix<double, 2, 3> J_pi;
        J_pi << fx * zc_inv, 0, -fx*lc(0)*zc_inv2,
                0, fy * zc_inv, -fy*lc(1)*zc_inv2;

        //std::cout << " J_pi ===== : " << J_pi << std::endl;

        Eigen::Vector3d pi(pose_i[0], pose_i[1], 0.0);
            
        //std::cout << " pi ===== : " << pi.transpose() << std::endl;

        Eigen::Matrix<double, 2, 3> J_pi_Rcw = J_pi * Rcw;

        Eigen::Matrix2d J_rotxy = (J_pi_Rcw * SE2::skew(lw - pi)).block<2,2>(0,0);
        Eigen::Matrix<double,2,1> J_z = -J_pi_Rcw.block<2,1>(0,2);

        Eigen::Matrix2d Sigma_all = Sigma_rotxy * J_rotxy * J_rotxy.transpose() +
                                    Sigma_z * J_z * J_z.transpose() + Sigma_u;

        //std::cout << " Sigma_all ===== : " << std::endl << Sigma_all << std::endl;

        Eigen::Matrix2d keypoint_info = Sigma_all.inverse(); 
        Eigen::Matrix2d sqrt_info = Eigen::LLT<Eigen::Matrix2d>(keypoint_info).matrixL().transpose();

        // residual
        Eigen::Map<Eigen::Matrix<double, 2, 1>> residual(residuals);

        residual = CamProject(lc) - Eigen::Vector2d(undist_keypt); 
        
        // LOG(INFO) << " CamProject(lc) ===== : " << CamProject(lc).transpose() << std::endl;
        // LOG(INFO) << " undist_keypt ===== : " << undist_keypt.transpose() << std::endl;

        // LOG(INFO) << " residual ===== : " << residual.transpose() << std::endl;

        if (jacobians) {
            Eigen::Matrix<double, 2, 3> J_pi;
            J_pi << fx * zc_inv, 0, -fx * lc(0) * zc_inv2,
                    0, fy * zc_inv, -fy * lc(1) * zc_inv2;

            Eigen::Matrix<double, 2, 3> J_pi_Rcw = J_pi * Rcw;

            if(jacobians[0]) { // dR
                Eigen::Map<Eigen::Matrix<double, 2, 1, Eigen::ColMajor>> jacobianR(jacobians[0]);
                jacobianR = (J_pi_Rcw * SE2::skew(lw - pi)).block<2,1>(0,2);
                jacobianR = sqrt_info * jacobianR;
            }
            if(jacobians[1]) { // dpose
                Eigen::Map<Eigen::Matrix<double, 2, 2, Eigen::RowMajor>> jacobian_pose(jacobians[1]);
                jacobian_pose = -J_pi_Rcw.block<2,2>(0,0);
                jacobian_pose = sqrt_info * jacobian_pose;
            }

            if(jacobians[2]) { // dpoint
                Eigen::Map<Eigen::Matrix<double, 2, 3, Eigen::RowMajor>> jacobian_point(jacobians[2]);
                jacobian_point = J_pi_Rcw;
                jacobian_point = sqrt_info * jacobian_point;

            }
        }    
        return true;
    }
private:
    Vector2d undistorted_point_;
    Eigen::Matrix3d intrinsic_matrix_;
    Eigen::Matrix4d extrinsic_matrix_;
};

}
#endif